import time
from bot.enterprise_wechat_gpt.config.conf import *
from bot.enterprise_wechat_gpt.util.tair_util import tair_util
from bot.enterprise_wechat_gpt.util.thread_local_util import local_data


def request_limit(url):
    rate_limit = tair_util.get("QW:RateLimit")
    if rate_limit is None:
        rate_limit = 15
        tair_util.set("QW:RateLimit", rate_limit)
    else:
        rate_limit = int(rate_limit)
    expire_key = "QW:RateLimit:" + url + ":Expire"
    incr_key = "QW:RateLimit:" + url
    absent = tair_util.set_if_absent(expire_key, 1)
    if absent:
        logging.info(f"url:{url} 一分钟内第一次请求")
        tair_util.set(incr_key, "1")
        if rate_limit < 1:
            logging.warning(f"url:{url} 限流数小于1, 直接限流")
            return True
        return False
    current_rate = tair_util.incr(incr_key)
    if current_rate > rate_limit:
        return True
    return False


def get_token(url):
    user_id = local_data.user_id
    retry_num = 12
    for i in range(retry_num):
        if not request_limit(url):
            return True
        logging.warning(f"userId:{user_id} url:{url} 一分钟内请求次数超过限制, 等待10s")
        time.sleep(10)
    else:
        # 连续尝试12次都被限流，返回True
        logging.info(f"userId:{user_id} url:{url} 连续尝试{retry_num}次都被限流, 直接放行")
        return True


if __name__ == '__main__':
    # 模拟多次请求
    for i in range(1, 20):
        token = get_token("test")
        if token:
            logging.info(f"Request {i} successful.")
        time.sleep(0.7)  # 模拟请求间隔